from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
import torch
from torch.utils.data import Dataset
import os

# 设置模型文件所在的本地路径
model_name = "C:/Users/86182/Desktop/sqsx_lqbz/model_download"  # 替换为BERT模型的本地路径

# 加载分词器和分类模型
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=3)  # 假设有3个评分等级

# 自定义数据集类，用于微调
class CustomDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: val[idx].clone().detach() for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

# 准备训练数据
train_texts = ["正确", "错误"]  # 替换为实际的训练文本
train_labels = [1, 2]  # 替换为实际的标签
encodings = tokenizer(train_texts, truncation=True, padding=True, return_tensors="pt")
dataset = CustomDataset(encodings, train_labels)

# 定义训练参数
training_args = TrainingArguments(
    output_dir='./results',            # 输出目录
    num_train_epochs=3,                # 训练的总epoch数
    per_device_train_batch_size=16,    # 每个设备的batch size
    per_device_eval_batch_size=64,     # 评估时的batch size
    warmup_steps=500,                  # 预热步数
    weight_decay=0.01,                 # 权重衰减
    logging_dir='./logs',              # 日志输出目录
    logging_steps=10,                  # 每隔多少步记录一次日志
    save_safetensors=False             # 禁用safetensors，使用传统保存方式
)

# 创建 Trainer 对象进行微调
trainer = Trainer(
    model=model,                       # 要微调的模型
    args=training_args,                # 训练参数
    train_dataset=dataset,             # 训练数据集
)

# 开始训练
trainer.train()

# 在保存之前确保所有张量是连续的
for name, param in model.named_parameters():
    if not param.is_contiguous():
        param.data = param.data.contiguous()

# 使用torch.save保存模型权重
output_dir = training_args.output_dir
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

model_path = os.path.join(output_dir, "pytorch_model.bin")
torch.save(model.state_dict(), model_path)

print(f"模型已保存至 {model_path}")
